fix: [BUG] qwen3_5 dense: mixed-precision NVFP4 crashes in ct_bf16_fuse (Float8 × BFloat16 promotion) - #275

Tested this approach out-of-tree — confirms it fixes #238 on a second checkpoint variant and a second platform.

I'm on a compiled build, so I couldn't apply the diff literally. Instead I patched one level down at ShardReader.get_tensor — the exact call this PR replaces — reusing the engine's own _dequant_fp8_weight for the bare-FP8 .weight case. Semantics should match the PR: verified byte-identical to a manual w.to(bf16) * scale, with weight_scale, weight_packed and plain BF16 tensors passing through untouched.

Env: freetoken 0.1.2+g932e87e8a, Python 3.12, Windows 11, sm_120, 16 GB, TP=1.
Checkpoint: a third-party Qwen3.8-27B multimodal NVFP4 export (Qwen3_5ForConditionalGeneration, compressed-tensors 0.18.0, format: "mixed-precision" — 232 float-quantized FP8 targets + 168 nvfp4-pack-quantized).

ft checkpoint was failing at ~2 s; it now completes:

wrote FTW checkpoint
  tensors: 875 weight + 0 experts_bank
  FTW: 30.06 GiB across 4 shard(s) (<= 8.0 GiB each)
  converted in 41.6s
[patch] dequantized 232 FP8 tensors: 8.71 GiB -> 17.42 GiB

ft serve pointed at the source dir: the Float8 promotion crash is gone. It now loads for ~38 s and stops on VRAM, which is expected for a dense 27B on a 16 GB card and unrelated to this PR:

Free memory before loading model: 14.62 GiB
OutOfMemoryError: CUDA out of memory. Tried to allocate 2.37 GiB.
GPU 0 has a total capacity of 15.89 GiB of which 0 bytes is free.
Of the allocated memory 23.13 GiB is allocated by PyTorch

Two caveats worth stating plainly:

• Numerical correctness is unverified. I could not run a token through the model on this hardware, so this confirms the crash is fixed and the pipeline completes — nothing more. Relatedly, the output carries 875 weight tensors where a pure-FP8 export of the same model produces 915. That is plausibly just different fusion between the two quant layouts, but I can't confirm it from here.
• Footprint. Dequantizing to bf16 takes those 232 tensors from 8.71 GiB to 17.42 GiB, so the resulting FTW (30.06 GiB) is larger than the 23 GiB source checkpoint. Keeping them FP8, as the quant_method: "fp8" path already does, would avoid that. Possibly out of scope for this fix, but it does make the mixed-precision path meaningfully more expensive than the pure-FP8 one.

One question: #238's traceback lands at loader.py:251 inside torch.cat, while mine lands at loader.py:256. That may just be line drift between the PyPI wheel and my dev build, but if ct_bf16_fuse genuinely has two Float8-reachable sites, it would be worth confirming this change covers both.

Correction to my report above — I spoke too soon. The conversion completing is not the end of the story: the resulting FTW does not load.

Serving it fails downstream, in model construction rather than in the weight walker:

File "freetoken/kernel/triton/nvfp4_linear.py", line 867, in load_state_dict
  s = state_dict.pop(_concat_prefix(prefix, "weight_scale"))
KeyError: 'model.layers.0.linear_attn.out_proj.weight_scale'

Root cause, and I don't think it is in scope for this PR: detect_compressed_tensors_nvfp4() in models/config.py returns True when any config_groups entry is NVFP4, and its docstring states the assumption it makes — "All targets: ["Linear"] are NVFP4 except the per-module ignore list". That does not hold for format: "mixed-precision". qwen3_5_moe/config.py then forces attn_quant = "nvfp4", so linear_attn.out_proj is built as Nvfp4DenseLinear and demands an NVFP4-form weight_scale. Dequantizing that tensor to bf16 (this PR's approach, and what I emulated) removes the scale entirely, so the KeyError is the direct consequence.

The deeper issue is that quantization is resolved into three global scalars (expert_quant / attn_quant / dense_quant), and this checkpoint cannot be described by any single assignment:

Fixing attn_quant alone would just move the failure to layers 56-63 via dense_quant.

So my earlier conclusion needs revising: this change removes the Float8 promotion crash, but on a mixed-precision checkpoint it produces an FTW the engine cannot consume. It converts, then dies at load.

Worth noting #208 appears to target exactly this, with per-module native storage reusing the existing Fp8PerTensor* kernels rather than dequantizing — and reports 21.8 GiB resident vs ~54 GB fully dequantized, with working generation. That also addresses the footprint point I raised above: the FP8 group here is per-row scaled, which Fp8PerTensorLinear already stores natively (weight fp8 + per-row fp32 weight_scale [out]), so the bf16 expansion is avoidable rather than inherent.

Is this PR still wanted as a narrower stopgap for checkpoints that only trip the loader path, or is it superseded by #208? Happy to test either.